import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from tqdm import tqdm


def to_float(x):
    return x.float()


if __name__ == '__main__':
    # 数据预处理
    data_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Lambda(to_float)
    ])

    # 导入数据
    train_data = datasets.ImageFolder(root='Data/train', transform=data_transform)
    test_data = datasets.ImageFolder(root='Data/test', transform=data_transform)

    train_loader = DataLoader(dataset=train_data, batch_size=50, shuffle=True, num_workers=4)
    test_loader = DataLoader(dataset=test_data, batch_size=50, num_workers=4)

    # 构建模型
    model = nn.Sequential(
        nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2),
        nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2),
        nn.Flatten(),
        nn.Linear(128 * 56 * 56, 1024),
        nn.ReLU(),
        nn.Linear(1024, 5)  # 5种绘画风格
    )

    device = torch.device('cpu')
    model = model.to(device)
    # device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
    # model = model.to(device)

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    criterion = criterion.to(device)  # 将损失函数移动到指定的设备上

    # 训练模型
    model.train()  # 启用模型训练模式
    for epoch in range(10):  # 迭代次数
        progress_bar = tqdm(enumerate(train_loader), total=len(train_loader))
        epoch_loss = 0
        correct = 0
        total = 0
        for i, (images, labels) in progress_bar:
            images = images.to(device)  # 将图像转移到与模型相同的设备上
            labels = labels.to(device).long()  # 将标签转移到与模型相同的设备上，并转换为LongTensor
            outputs = model(images)
            loss = criterion(outputs, labels)
            epoch_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            progress_bar.set_description(f'Epoch: {epoch + 1}, Step: {i + 1}, Loss: {loss.item()}')

        epoch_loss /= len(train_loader)
        epoch_accuracy = 100 * correct / total
        print(f'Epoch: {epoch + 1}, Loss: {epoch_loss}, Accuracy: {epoch_accuracy}%')

    # 保存模型参数
    torch.save(model.state_dict(), 'models/model.pth')

    # 加载模型参数
    model.load_state_dict(torch.load('models/model.pth'))

    # 测试模型
    model.eval()
    with torch.no_grad():
        correct = 0
        total = 0
        for images, labels in test_loader:
            # 将图像和标签都移动到指定的设备上
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        print(f'Accuracy: {100 * correct / total}%')
